# -*- coding: utf-8 -*-

import os
import io
from PIL import Image, TiffTags, ExifTags
import stat
Image.MAX_IMAGE_PIXELS = None
import time
import rawpy
import imageio
import subprocess
from tkinter import messagebox
import sys
from psd_tools import PSDImage 
from psd_tools.constants import Resource
from PIL import ImageCms

def rmtree(top):
    for root, dirs, files in os.walk(top, topdown=False):
        for name in files:
            filename = os.path.join(root, name)
            os.chmod(filename, stat.S_IWUSR)
            os.remove(filename)
        for name in dirs:
            os.rmdir(os.path.join(root, name))
    os.rmdir(top) 

# 删除文件或者文件夹
def Delete(path):
    if os.path.isdir(path):
        rmtree(path)
    else:
        os.chmod(path, stat.S_IWUSR)
        os.remove(path)
        

def ConvertToTiff(filePathOrigin, filePathTarget, deleteOriginAfterConvert = False):

    # tiff压缩参数 None, "group3", "group4", "jpeg", "lzma", "packbits", "tiff_adobe_deflate", "tiff_ccitt", "tiff_lzw", "tiff_raw_16", "tiff_sgilog", "tiff_sgilog24", "tiff_thunderscan", "webp", "zstd"
    compression = "tiff_lzw"
    TiffTags.LIBTIFF_CORE.add(317)
    try:
        img = Image.open(filePathOrigin)
        img.save(filePathTarget,
                tiffinfo={317: 2},
                compression = compression,
                icc_profile = img.info.get('icc_profile',''))
        img.close()
        if deleteOriginAfterConvert:
            Delete(filePathOrigin)
        return f"转换完成\t{filePathTarget}"
    except Exception as e:
        return f"ERROR!!! 转换失败\t{filePathTarget}\t{e}"
        
def CopyExifFrom(filePathOrigin, filePathTarget):
    cmd = f".\exiftool -overwrite_original -TagsFromFile \"{filePathOrigin}\" \"-all:all>all:all\" \"{filePathTarget}\" "
    result = subprocess.run(cmd)
    print(result)

def IsLocked(filepath):
    locked = None
    file_object = None
    if os.path.exists(filepath):
        try:
            buffer_size = 8
            # Opening file in append mode and read the first 8 characters.
            file_object = open(filepath, 'a', buffer_size)
            if file_object:
                locked = False
        except IOError as message:
            locked = True
        finally:
            if file_object:
                file_object.close()
    return locked
        
def ConvertToJpeg(filePathOrigin, filePathTarget, deleteOriginAfterConvert = False):
    # 目标文件后缀和目标格式格式
    targetTail = ''
    targetSuffix = '.jpg'
    # jpeg压缩参数
    quality = 85        
    subsampling = "4:2:0"
    # 是否转成sRGB
    convertToSRGB = False
    # 渐进式压缩
    progressive = True
    
    fileNameFull = os.path.basename(filePathOrigin)
    (fileNameLeft, suffix) = os.path.splitext(fileNameFull)
    (filePathLeft, suffix) = os.path.splitext(filePathOrigin)

    # 如果源文件名中已经有了附加后缀，则不处理
    if len(targetTail) > 0 and fileNameLeft.find(targetTail) != -1:
        return 
    filePathTarget = filePathLeft + targetTail + targetSuffix
    

    try:
        if IsLocked(filePathOrigin):
            return f"ERROR!!! 文件被锁定\t{filePathOrigin}"
        if suffix.lower() in ['.nef','.dng']:
            # RAW DATA
            raw = rawpy.imread(filePathOrigin)
            with rawpy.imread(filePathOrigin) as raw:
                if convertToSRGB:
                    rgb = raw.postprocess(use_camera_wb=True, use_auto_wb=False, user_flip=0, output_color=rawpy.ColorSpace.sRGB)
                else:
                    rgb = raw.postprocess(use_camera_wb=True, use_auto_wb=False, user_flip=0)
                imageio.imsave(filePathTarget, rgb)
                # 从源文件复制EXIF信息
                CopyExifFrom(filePathOrigin, filePathTarget)
        else:
            # 图片
            if suffix.lower() in ['.psd','.psb']:
                # Photoshop文件
                psd = PSDImage.open(filePathOrigin)
                img = psd.composite()
                # icc字节对象
                iccBytesSource = psd.image_resources.get_data(Resource.ICC_PROFILE)
            else:
                # 其他图片格式
                img = Image.open(filePathOrigin)
                # icc字节对象
                iccBytesSource = img.info.get('icc_profile','')
            if iccBytesSource and len(iccBytesSource) > 0:
                imgProfileSource = ImageCms.ImageCmsProfile(io.BytesIO(iccBytesSource))
                imageProfileSourceName = ImageCms.getProfileName(imgProfileSource)

            if convertToSRGB and imageProfileSourceName and 'sRGB' not in imageProfileSourceName:
                # 源文件不是sRGB，并且需要转成sRGB，进行色彩空间转换操作
                imgProfileTarget = ImageCms.ImageCmsProfile(ImageCms.createProfile('sRGB'))
                img = ImageCms.profileToProfile(img, imgProfileSource, imgProfileTarget, 0, "RGB")
                iccFileTarget = imgProfileTarget.tobytes()
            else:
                iccFileTarget = iccBytesSource
                    
            if img.format == 'JPEG':
                # 如果源文件是JPEG，保留源subsampling
                subsampling = "keep"
            
            (w,h) = img.size
            if w > 65500 or h > 65500:
                # 如果源文件长边超过65500，则缩小
                if w > h:
                    # 长边是宽
                    img = img.resize((65500, int(65500 / w * h))) 
                else:
                    # 长边是高
                    img = img.resize(int(65500 / h * w), 65500)
                
            
            img.save(filePathTarget,
                    quality = quality,
                    subsampling = subsampling,
                    progressive = progressive,
                    icc_profile = iccFileTarget)
            img.close()

        if deleteOriginAfterConvert:
            Delete(filePathOrigin)
        return f"转换完成\t{filePathOrigin}"
    except Exception as e:
        return f"ERROR!!!\t{filePathOrigin}\t{e} "
    


def main():
    T1 = time.perf_counter()
    # 需要转换的目录
    base = r"E:\deleteme\标准测试图案"
    # 需要转换的格式
    suffixToConvert = {'.bmp', '.tif', '.tiff', '.psd', '.psb', '.nef'}
    #suffixToConvert = {'.jpg'}
    # 目标文件后缀和目标格式格式
    targetTail = ''
    targetSuffix = '.jpg'
    # 是否删除源文件
    deleteOriginAfterConvert = False
    
    # 待处理的文件
    listFileToProcess = []
    strOverWrite = ''
    for root, dirs, files in os.walk(base, topdown=False):
        for fileName in files:
            filePathOrigin = os.path.join(root, fileName)
            # 判断后缀名是不是在转换名单里
            fileNameFull = os.path.basename(filePathOrigin)
            (fileNameLeft, suffix) = os.path.splitext(fileNameFull)
            (filePathLeft, suffix) = os.path.splitext(filePathOrigin)
            
            # 不处理无文件名的文件，比如.DS_STORE，._山磨水溪
            if fileNameFull[0] == '.':
                continue
            if fileNameLeft is None:
                continue
            if len(fileNameLeft) == 0:
                continue
            # 如果源文件名中已经有了附加后缀，则不处理
            if len(targetTail) > 0 and fileNameLeft.find(targetTail) != -1:
                continue 
            # 源文件不是待处理格式的话不处理
            if suffix.lower() not in suffixToConvert:
                continue
            # 组成目标文件名
            filePathTarget = filePathLeft + targetTail + targetSuffix
            
            # 添加到待处理
            listFileToProcess.append((filePathOrigin, filePathTarget))
            if filePathOrigin == filePathTarget:
                # 目标文件和源文件名相同，会覆盖
                strOverWrite += f"{filePathOrigin}\n"

    # 会删除原文件，警告用户
    if deleteOriginAfterConvert:
        if "no" == messagebox.askquestion(f"转换任务{base}", f"已开启转换后删除原文件，是否继续？"):
            print('----------------------------- 用户选择不继续，运行时间:%s秒' % (time.perf_counter() - T1))
            sys.exit()
    # 会产生覆盖，警告用户
    if len(strOverWrite) > 0:
        if "no" == messagebox.askquestion(f"转换任务{base}", f"下列文件转换后会覆盖掉源文件，是否继续？\n{strOverWrite}"):
            print('----------------------------- 用户选择不继续，运行时间:%s秒' % (time.perf_counter() - T1))
            sys.exit()
    i = 0
    count = len(listFileToProcess)
    
    for (filePathOrigin, filePathTarget) in listFileToProcess:
        i = i+1

        fileNameFull = os.path.basename(filePathOrigin)
        (fileNameLeft, suffix) = os.path.splitext(fileNameFull)
        (filePathLeft, suffix) = os.path.splitext(filePathOrigin)
        if targetSuffix == ".jpg":
            result =  ConvertToJpeg(filePathOrigin, filePathTarget, deleteOriginAfterConvert)
        elif targetSuffix == ".tif":
            result =  ConvertToTiff(filePathOrigin, filePathTarget, deleteOriginAfterConvert)
        else:
            msg = f"{i}/{count}     不支持的目标格式{targetSuffix}"
            continue
        msg = f"{i}/{count}     {result}"
        print(msg)

    print('----------------------------- 程序运行结束，运行时间:%s秒' % (time.perf_counter() - T1))
    
if __name__ == '__main__':
    main()
    